import oss2
import os
import argparse
from PIL import Image
from datetime import datetime
import logging

def setup_logger(name, log_file, level=logging.INFO):
    """Function to setup logging"""
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler = logging.FileHandler(log_file)
    handler.setFormatter(formatter)
    logger = logging.getLogger(name)
    logger.setLevel(level)
    logger.addHandler(handler)
    return logger

def main(args):
    auth = oss2.Auth(args.access_key_id, args.access_key_secret)
    bucket = oss2.Bucket(auth, 'oss-cn-shenzhen.aliyuncs.com', args.bucket_name)

    # 转换输入的起止日期为datetime格式
    start_date = datetime.strptime(args.start_date, "%Y%m%d")
    end_date = datetime.strptime(args.end_date, "%Y%m%d")

    log_folder = os.path.join(args.local_path, 'oss_download_log')
    if not os.path.exists(log_folder):
        os.makedirs(log_folder)

    log_file_name = f'oss_download_{datetime.now().strftime("%Y%m%d")}.log'
    log_file = os.path.join(log_folder, log_file_name)
    logger = setup_logger(__name__, log_file)

    for sn in args.sn_list:
        prefix = sn + '/'
        for obj in oss2.ObjectIterator(bucket, prefix=prefix):
            if obj.key[-1] != "/":  # 确保是文件而不是文件夹

                # 获取日期部分并转换为 datetime 对象
                date_str = obj.key.rsplit('/', 2)[1][2:10]  # OD2023111255484948 -> 20231112
                try:
                    obj_date = datetime.strptime(date_str, "%Y%m%d")
                except ValueError:  # 如果日期格式不正确，跳过
                    continue

                # 判断日期是否在指定区间内
                if start_date <= obj_date <= end_date:
                    logger.info(f'开始下载SN {sn}的文件：{obj.key}，日期范围：{args.start_date} - {args.end_date}')
                    
                    print('开始下载文件：', obj.key)

                    local_file_path = os.path.join(args.local_path, obj.key.rsplit('/', 1)[0])
                    os.makedirs(local_file_path, exist_ok=True)
                    new_file_name = obj.key.rsplit('/', 2)[0] + '_' + obj.key.rsplit('/', 2)[1] + '_' + obj.key.rsplit('/', 1)[1]
                    bucket.get_object_to_file(obj.key, os.path.join(local_file_path, new_file_name))
                    

                    # 对下载的文件进行 resize
                    image_path = os.path.join(local_file_path, new_file_name)
                    image = Image.open(image_path)
                    image = image.resize((960, 540))
                    image.save(os.path.join(local_file_path, new_file_name), quality=95)
                # else:  # 可以在此处添加日志记录或其他操作，以便了解跳过了哪些文件


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='批量下载阿里云OSS文件夹中的文件')
    parser.add_argument('--access_key_id', type=str, default='LTAI5tSGrQNwcFQKfBikvGz8', help='你的AccessKeyId')
    parser.add_argument('--access_key_secret', type=str, default='sMgtMLNhRqfsmDkABgRcnumQvNlSQB', help='你的AccessKeySecret')
    parser.add_argument('--bucket_name', type=str, default='daqi-retail', help='你的BucketName')
    parser.add_argument('--sn_list', type=str, nargs='+', default=['SN001', 'SN005'], help='输入SN号list，例如：SN001 SN002 SN003')
    parser.add_argument('--local_path', type=str, default='daqi_retail', help='本地文件路径')
    # 增加起止日期参数，格式为 YYYYMMDD
    parser.add_argument('--start_date', type=str, default='20240410',  help='起始日期 (YYYYMMDD)')
    parser.add_argument('--end_date', type=str, default='20240529',  help='结束日期 (YYYYMMDD)')

    args = parser.parse_args()

    main(args)